#include "imuprocessor.h"

#define DEBUG false

double square(const double &arg)
{
  return arg*arg;
}
// helper function to print a matrix to Serial
template <int N, int M>
void printmatrix(const MatrixNxM<N,M> &m)
{
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < M; j++)
    {
      printf("%.6f ", m(i, j));
    }
    printf("\n");
  }
}

// helper function to print a vector to Serial
template <int N>
void printvector(const VectorND<N> &v)
{
  for (int i = 0; i < N; i++)
  {
    printf("%.6f ", v(i));
  }
  printf("\n");
}

IMUProcessor::IMUProcessor()
{
  id = 0;
  prev_id = 1;


}

void IMUProcessor::setMeasurementErrors(const Vector3D &a, const Vector3D &m, const Vector3D &g)
{
  m_aerr = a;
  m_merr = m;
  m_gerr = g;
}
void IMUProcessor::setInitialValues(const Vector3D &a, const Vector3D &m, const Vector3D &g)
{
  Quaternion qpos, qposerr;

  vectorToQuaternion(a, Vector3D(), qpos, qposerr);

  for (int i = 0; i < 4; i++)
    pos[id](i) = qpos(i);

  for (int i = 0; i < 3; i++)
    pos[id](4+i) = g(i);




  // initial covariance 
  // TODO: errors on bias offset? 
  P[id](0,0) = 0.1*0.1;
  P[id](1,1) = 0.1*0.1;
  P[id](2,2) = 0.1*0.1;
  P[id](3,3) = 0.1*0.1;
  for (int i = 0; i < 3; i++)
    P[id](4+i,4+1) = 0.01*0.01;



  // set references
  rm = m;
  rm /= rm.length();
  // the acceleration is taken to (0,0,1)






}

void IMUProcessor::vectorToQuaternion(const Vector3D &a, const Vector3D &aerr, Quaternion &qpos, Quaternion &qposerr)
{
  qposerr = Quaternion();

  double pitch, roll;

  Vector3D acc = a;
  acc /= acc.length();

  pitch = asin(acc.x());
  roll = atan2(-acc.y(),acc.z());

  // all errors are NOT squared
  double pitcherr1, pitcherr2, pitcherr;
  double rollerr1, rollerr2, rollerr;

  pitcherr1 = aerr.x() / fabs(acc.y()+acc.y() + acc.z()*acc.z());
  pitcherr2 = (acc.y()*acc.y()*aerr.y() + acc.z()*acc.z()*aerr.z())/acc.x()/(acc.y()*acc.y() + acc.z()*acc.z());
  pitcherr = pitcherr1;
  if (pitcherr2 < pitcherr1)
    pitcherr = pitcherr2;

  rollerr1 = sqrt(aerr.y() + aerr.z()*acc.y()*acc.y()/acc.z()/acc.z()) / ( (1+acc.y()*acc.y()/acc.z()/acc.z()) /acc.z()) ;
  rollerr1 *= rollerr1;
  rollerr2 = sqrt(aerr.y() + pitcherr*pitcherr*acc.x()*acc.x()*acc.y()*acc.y()/(acc.y()*acc.y()+acc.z()*acc.z()))/acc.z();
  rollerr2 *= rollerr2;
  rollerr = rollerr1;
  if (rollerr2 < rollerr1)
    rollerr = rollerr2;

  double cp, sp, cr, sr;
  cp = cos(pitch/2);
  sp = sin(pitch/2);
  cr = cos(roll/2);
  sr = sin(roll/2);

  qpos = Quaternion(
    cp*cr,
    cp*sr,
    sp*cr,
    sp*sr
  );
   
  cp *= cp;
  sp *= sp;
  cr *= cr;
  sr *= sr;

  // qpos err IS squared
  qposerr = Quaternion(
    pitcherr * cr * sp + rollerr  * cp * sr,
    rollerr  * cp * cr + pitcherr * sp * sr,
    pitcherr * cp * cr + rollerr  * sp * sr,
    rollerr  * cr * sp + pitcherr * cp * sr
  );  
  qposerr *= 0.25;
}

void IMUProcessor::eliminateYaw(Quaternion &q)
{
  double pitch, roll, yaw;

  roll = 2* (q(0) * q(1) - q(2) * q(3)) / (q(0)*q(0) - q(1)*q(1) - q(2)*q(2) + q(3)*q(3));
  roll = atan(roll);
  
  pitch = 2*(q(0)*q(2) + q(3)*q(1));
  pitch = asin(pitch);

  yaw = 2* (q(0) * q(3) - q(1) * q(2)) / (q(0)*q(0) + q(1)*q(1) - q(2)*q(2) - q(3)*q(3));
  yaw = atan(yaw);

  double cp, sp, cr, sr;
  cp = cos(pitch/2);
  sp = sin(pitch/2);
  cr = cos(roll/2);
  sr = sin(roll/2);

  q = Quaternion(
    cp*cr,
    cp*sr,
    sp*cr,
    sp*sr
  );
}


void IMUProcessor::estimatePositionQuaternionNewton(/* outputs*/ Quaternion &angle, Quaternion &angleerr)
{
  angleerr = Quaternion(0,0,0,0);

  // compute best guess to seed the Gauss-Newton search
  //  compute change from bias-corrected gyroscope data
  Quaternion prevpos = Quaternion(pos[prev_id](0), pos[prev_id](1), pos[prev_id](2), pos[prev_id](3));
  Quaternion dposgyro = -(double)dt / 1000 / 1000 * 0.5 * (prevpos * Quaternion(0, g.x() - pos[prev_id](4), g.y() - pos[prev_id](5), g.z() - pos[prev_id](6)));
  //  estimate new position:
  Quaternion q(pos[prev_id](0), pos[prev_id](1), pos[prev_id](2), pos[prev_id](3));
  q += dposgyro;

  // normalize
  q /= q.length();

  q = Quaternion(1,0,0,0);

  //printf("q-init: %f %f %f %f", q(0), q(1), q(2), q(3));

  Vector3D aa, mm;
  aa = a;
  aa /= aa.length();
  mm = m;
  mm /= mm.length();


  // proceed with the Gauss-Newton search

  MatrixNxM<6,4> J;
  VectorND<6> esterr;
  VectorND<4> step;
  Matrix4x4 JJ;

  // precompute commonly occuring products
  double q0q0, q0q1, q0q2, q0q3, q1q1, q1q2, q1q3, q2q2, q2q3, q3q3; 
  double r1r1, r1r2, r1r3, r2r2, r2r3, r3r3; 
  r1r1 = rm.x()*rm.x();
  r1r2 = rm.x()*rm.y();
  r1r3 = rm.x()*rm.z();
  r2r2 = rm.y()*rm.y();
  r2r3 = rm.y()*rm.z();
  r3r3 = rm.z()*rm.z();

  for (int iteration = 0; iteration < 4; iteration++)
  {
    // tack on a factor of 2 to q (!!!)
    q *= 2;


    J(0,0) =  q(2);
    J(1,0) = -q(1);
    J(2,0) =  q(0);
    J(3,0) =  rm.x()*q(0) + rm.z()*q(2) - rm.y()*q(3);
    J(4,0) =  rm.y()*q(0) - rm.z()*q(1) + rm.x()*q(3);
    J(5,0) =  rm.z()*q(0) + rm.y()*q(1) - rm.x()*q(2);
    J(0,1) =  q(3);
    J(1,1) = -q(0);
    J(2,1) = -q(1);
    J(3,1) =  rm.x()*q(1) + rm.y()*q(2) + rm.z()*q(3);
    J(4,1) = -J(5,0);
    J(5,1) =  J(4,0);
    J(0,2) =  q(0);
    J(1,2) =  q(3);
    J(2,2) = -q(2);
    J(3,2) =  J(5,0);
    J(4,2) =  J(3,1);
    J(5,2) = -J(3,0);
    J(0,3) =  q(1);
    J(1,3) =  q(2);
    J(2,3) =  q(3);
    J(3,3) = -J(4,0);
    J(4,3) =  J(3,0);
    J(5,3) =  J(3,1);
    // un-tack on a factor of 2 to q (!!!)
    q /= 2;


    q0q0 = q(0)*q(0);
    q0q1 = q(0)*q(1);
    q0q2 = q(0)*q(2);
    q0q3 = q(0)*q(3);
    q1q1 = q(1)*q(1);
    q1q2 = q(1)*q(2);
    q1q3 = q(1)*q(3);
    q2q2 = q(2)*q(2);
    q2q3 = q(2)*q(3);
    q3q3 = q(3)*q(3);

    esterr(0) = aa.x() - 2*( q0q2 + q1q3);
    esterr(1) = aa.y() - 2*(-q0q1 + q2q3);
    esterr(2) = aa.z() - q0q0 + q1q1 + q2q2 - q3q3;
    esterr(3) = mm.x() - 2*rm.y()*( q1q2 - q0q3) - 2*rm.z()*( q0q2 + q1q3) - rm.x()*(q0q0 + q1q1 - q2q2 - q3q3);
    esterr(4) = mm.y() - 2*rm.x()*( q1q2 + q0q3) - 2*rm.z()*(-q0q1 + q2q3) - rm.y()*(q0q0 - q1q1 + q2q2 - q3q3);
    esterr(5) = mm.z() - 2*rm.x()*(-q(0)*q(2) + q(1)*q(3)) - 2*rm.y()*( q(0)*q(1) + q(2)*q(3)) - rm.z()*(q(0)*q(0) - q(1)*q(1) - q(2)*q(2) + q(3)*q(3));

    double ea = 0.1, em = 1;
    esterr(0) *= ea;
    esterr(1) *= ea;
    esterr(2) *= ea;
    esterr(3) *= em;
    esterr(4) *= em;
    esterr(5) *= em;

    JJ = J.transpose()*J; 
    //printmatrix(JJ);
    JJ = JJ.inverse();
    //printmatrix(JJ);
    step = JJ*((J.transpose())*esterr);

    //double steplen = step.length();

    q = q + step;

  }
  q /= q.length();
  angle = q;

  //  ok this is just riddiculous
  angleerr(0) = 2/((8*q1q1)/m_aerr.y() + (8*q2q2)/
     m_aerr.x() - (4*(a.z() - 3*q0q0 + q1q1 + q2q2 - q3q3))/m_aerr.z() + 
        (1/m_merr.x())*(-4*rm.x()*(m.x() + (-3*q0q0 - q1q1 + q2q2 + q3q3)*rm.x()) + 
       8*(q1q2 - 3*q0q3)*r1r2 + 
             8*q3q3*r2r2 + 
       8*(3*q0q2*rm.x()+ q1q3*rm.x()- 2*q2q3*rm.y())*rm.z()+ 8*q2q2*r3r3) + 
        (8*square(q(2)*rm.x()- q(1)*rm.y()) - 4*m.z()*rm.z()+ 
       8*(-3*q0q2*rm.x()+ q1q3*rm.x()+ 3*q0q1*rm.y()+ q2q3*rm.y())*rm.z()- 
             4*(-3*q0q0 + q1q1 + q2q2 - q3q3)*r3r3)/m_merr.z() + 
        (1/
       m_merr.y())*(4*(rm.y()*(-m.y() + 2*q1q2*rm.x()+ 3*q0q0*rm.y()- q1q1*rm.y()+ 
            q2q2*rm.y()) + 
                q3q3*(2*r1r1 - r2r2) - 6*q0q1*r2r3 + 2*q1q1*r3r3 + 
                2*q(3)*(3*q(0)*r1r2 - 2*q(1)*r1r3 + q(2)*r2r3))));

  angleerr(1) =  1/(2*((2*q0q0)/m_aerr.y() + (2*q3q3)/
       m_aerr.x() + (a.z() - q0q0 + 3*q1q1 + q2q2 - q3q3)/m_aerr.z() + 
           (2*square(q(3)*rm.x()+ q(0)*rm.y()) + (m.z() + 2*q0q2*rm.x()- 6*q1q3*rm.x()- 
            6*q0q1*rm.y()- 2*q2q3*rm.y())*rm.z()- 
                (q0q0 - 3*q1q1 - q2q2 + q3q3)*r3r3)/m_merr.z() + 
           (1/m_merr.x())*((-m.x())*rm.x()+ q0q0*r1r1 + 3*q1q1*r1r1 - 
         q2q2*r1r1 - q3q3*r1r1 + 6*q1q2*r1r2 + 
                2*q2q2*r2r2 + 6*q1q3*r1r3 + 4*q2q3*r2r3 + 
         2*q3q3*r3r3 + 2*q(0)*rm.x()*((-q(3))*rm.y()+ q(2)*rm.z())) + 
           (1/
         m_merr.y())*(rm.y()*(m.y() - 2*q0q3*rm.x()- q0q0*rm.y()+ 3*q1q1*rm.y()+ 
            q3q3*rm.y()) + q2q2*(2*r1r1 - r2r2) + 
                6*q0q1*r2r3 + 2*q0q0*r3r3 - 
         2*q(2)*(3*q(1)*r1r2 + 2*q(0)*r1r3 + q(3)*r2r3))));

  angleerr(2) =  1/(2*((2*q0q0)/m_aerr.x() + (2*q3q3)/
       m_aerr.y() + (a.z() - q0q0 + q1q1 + 3*q2q2 - q3q3)/m_aerr.z() + 
           (2*square(q(0)*rm.x()- q(3)*rm.y()) + (m.z() - 
            2*(-3*q0q2*rm.x()+ q1q3*rm.x()+ q0q1*rm.y()+ 3*q2q3*rm.y()))*rm.z()- 
                (q0q0 - q1q1 - 3*q2q2 + q3q3)*r3r3)/m_merr.z() + 
           (1/
         m_merr.y())*(rm.y()*(-m.y() + 
            2*q0q3*rm.x()+ (q0q0 + 3*q2q2 - q3q3)*rm.y()) + 
         q1q1*(2*r1r1 - r2r2) + 
                6*q2q3*r2r3 + 2*q3q3*r3r3 + 
         q(1)*(6*q(2)*r1r2 + 4*q(3)*r1r3 - 2*q(0)*r2r3)) + 
           (1/m_merr.x())*(m.x()*rm.x()- q1q1*r1r1 + 3*q2q2*r1r1 + q3q3*r1r1 - 
         6*q1q2*r1r2 + 2*q1q1*r2r2 - 
                2*q1q3*r1r3 + 
         2*q(0)*(q(3)*r1r2 - 3*q(2)*r1r3 + 2*q(1)*r2r3) - 
         q0q0*(r1r1 - 2*r3r3))));


  angleerr(3) =  2/((8*q1q1)/m_aerr.x() + (8*q2q2)/
     m_aerr.y() - (4*(a.z() - q0q0 + q1q1 + q2q2 - 3*q3q3))/m_aerr.z() + 
        (8*square(q(1)*rm.x() + q(2)*rm.y()) - 
       4*(m.z() + 2*q0q2*rm.x() - 6*q1q3*rm.x() - 2*q0q1*rm.y() - 6*q2q3*rm.y())*rm.z() - 
             4*(-q0q0 + q1q1 + q2q2 - 3*q3q3)*r3r3)/m_merr.z() + 
        (1/
       m_merr.x())*(4*(m.x()*rm.x() + 
         rm.x()*((-q1q1 + q2q2 + 3*q3q3)*rm.x() - 2*q1q2*rm.y()) - 
                q0q0*(r1r1 - 2*r2r2) - 6*q1q3*r1r3 + 2*q1q1*r3r3 + 
                q(0)*(6*q(3)*r1r2 - 2*q(2)*r1r3 - 4*q(1)*r2r3))) + 
        (1/
       m_merr.y())*(4*(rm.y()*(m.y() - 
            2*q1q2*rm.x() + (q1q1 - q2q2 + 3*q3q3)*rm.y()) + 
         q0q0*(2*r1r1 - r2r2) - 
                6*q2q3*r2r3 + 2*q2q2*r3r3 + 
         q(0)*(-6*q(3)*r1r2 + 4*q(2)*r1r3 + 2*q(1)*r2r3))));

}





void IMUProcessor::process(const Vector3D &aa, const Vector3D &mm, const Vector3D &gg, const int dt)
{
  // update array index first
  prev_id = id;
  id++;
  id = id % 2;

  // store the various readings:
  a = aa;
  m = mm;
  g = gg;


  // estimate the "raw" angle from acceleration only
  //   this is completely useless but is a good reference for comparing how the second method does
  vectorToQuaternion(a, Vector3D(), posrawacc[id], posrawaccerr[id]);



  // estimate the angle used as input for Kalman filter. 
  estimatePositionQuaternionNewton(/* outputs*/ posnwt[id], posnwterr[id]);

  /*
  printf("previous position ");
  printvector(pos[prev_id]);
  //printf("   error: ");
  //printvector(poserr[prev_id]);
  printf("\ncurrent newton estimate ");
  printvector(posnwt[id]);
  printf("   error: ");
  printvector(posnwterr[id]);
  printf("\n");
  */


  //posnwt[id] = posrawacc[id];
  //posnwterr[id] = posrawaccerr[id];

  // kalman filtering
  // *************
  // KALMAN FILTER
  // *************

  // Let me first tell you how this is NOT done
  // x_n = A.x_{n-1} + B.u_n
  // x_n -- 4 quaternion components
  // u_n -- controls, (0, w_x, w_y, w_z) gyro reading
  // A   -- just the identity
  // B   -- quaternion multiplication: B.u_n = - t * 0.5 x_{n-1} ** u_n 
  //    in this way B is a function of x_n -- noise / variance.
  // 
  // use extended kalman filter approach instead:
  //   x_n = f(x_{n-1}, u_n)
  // so that after linearization:
  //   x_n = A.x_{n-1}
  // where A encodes the propagation of old x_{n-1} AND the quaternion multiplication by the gyro readings
  // this allows for reliable Q estimation



  double deltat = 0.5 * (double)dt/1000.0/1000.0; // tack the 0.5 factor onto t to reduce mult.

  MatrixNxN<7> A = MatrixNxN<7>::identity();
  A(0, 1) =  g.x() * deltat;
  A(0, 2) =  g.y() * deltat;
  A(0, 3) =  g.z() * deltat;

  A(1, 0) = -g.x() * deltat;
  A(1, 2) = -g.z() * deltat;
  A(1, 3) =  g.y() * deltat;

  A(2, 0) = -g.y() * deltat;
  A(2, 1) =  g.z() * deltat;
  A(2, 3) = -g.x() * deltat;

  A(3, 0) = -g.z() * deltat;
  A(3, 1) = -g.y() * deltat;
  A(3, 2) =  g.x() * deltat;

  A(0, 4) = -deltat * pos[prev_id](1);
  A(0, 5) = -deltat * pos[prev_id](2);
  A(0, 6) = -deltat * pos[prev_id](3);

  A(1, 4) =  deltat * pos[prev_id](0);
  A(1, 5) = -deltat * pos[prev_id](3);
  A(1, 6) =  deltat * pos[prev_id](2);

  A(2, 4) =  deltat * pos[prev_id](3);
  A(2, 5) =  deltat * pos[prev_id](0);
  A(2, 6) = -deltat * pos[prev_id](1);

  A(3, 4) = -deltat * pos[prev_id](2);
  A(3, 5) =  deltat * pos[prev_id](1);
  A(3, 6) =  deltat * pos[prev_id](0);

  double sx2, sy2, sz2;  // gyro sigmas sq
  sx2 = m_gerr.x();
  sy2 = m_gerr.y();
  sz2 = m_gerr.z();
  //sx2 = sy2 = sz2 = 1;


  double factor = 0.4 * (double)dt/1000.0/1000.0;
  sx2 *= factor*factor;
  sy2 *= factor*factor;
  sz2 *= factor*factor;

  double qw, qx, qy, qz; // tack on the factor of 0.5 here 
  qw = pos[prev_id](0) * 0.5;
  qx = pos[prev_id](1) * 0.5;
  qy = pos[prev_id](2) * 0.5;
  qz = pos[prev_id](3) * 0.5;
  
  MatrixNxN<7> Q;
  Q(0,0) =  qx*qx*sx2 + qy*qy*sy2 + qz*qz*sz2;
  Q(0,1) = -qw*qx*sx2 - qy*qz*sy2 + qy*qz*sz2;
  Q(0,2) =  qx*qz*sx2 - qw*qy*sy2 - qx*qz*sz2;
  Q(0,3) = -qx*qy*sx2 + qx*qy*sy2 - qw*qz*sz2;

  Q(1,0) = Q(0,1);
  Q(1,1) =  qw*qw*sx2 + qz*qz*sy2 + qy*qy*sz2;
  Q(1,2) = -qw*qz*sx2 + qw*qz*sy2 - qx*qy*sz2;
  Q(1,3) =  qw*qy*sx2 - qx*qz*sy2 - qw*qy*sz2;
 
  Q(2,0) = Q(0,2);
  Q(2,1) = Q(1,2);
  Q(2,2) =  qz*qz*sx2 + qw*qw*sy2 + qx*qx*sz2;
  Q(2,3) = -qy*qz*sx2 - qw*qx*sy2 + qw*qx*sz2;

  Q(3,0) = Q(0,3);
  Q(3,1) = Q(1,3);
  Q(3,2) = Q(2,3);
  Q(3,3) =  qy*qy*sx2 + qx*qx*sy2 + qw*qw*sz2;

  double sigbias = 0.000000001;
  Q(4,4) = sigbias * sigbias;
  Q(5,5) = sigbias * sigbias;
  Q(6,6) = sigbias * sigbias;

  

  // compute posrawgyr on the way
  {
    VectorND<7> temp = pos[prev_id];
    for (int i = 0; i < 4; i++)
      temp(i) = posrawgyr[prev_id](i);

    temp = A * temp;

    for (int i = 0; i < 4; i++)
    {
      posrawgyr[id](i) = temp(i);
      posrawgyrerr[id](i) = 0;
    }
    posrawgyr[id] /= posrawgyr[id].length();
  }
  

  // prediction step
  pos[id] = A * pos[prev_id];
  // cheat and normalize;
  double quatlength = 0;
  for (int i = 0; i < 4; i++)
    quatlength += pos[id](i)*pos[id](i);
  quatlength = sqrt(quatlength);
  for (int i = 0; i < 4; i++)
    pos[id](i) /= quatlength;
  // prediction covariance 
  P[id] = A * P[prev_id] * A.transpose() + Q;


  if (DEBUG)
  {
    printf("** BEFORE UPDATE **\n");
    printf("matrix A: \n");
    printmatrix(A);
    printf("matrix Q:\n");
    printmatrix(Q);
    printf("matrix P:\n");
    printmatrix(P[id]);
  }

  // update step
  Matrix4x4 R = Matrix4x4::fromVector(posnwterr[id]);
  for (int i = 0; i < 4; i++)
    R(i,i) *= R(i,i);

  MatrixNxM<4,7> H;
  for (int i = 0; i < 4; i++)
    H(i,i) = 1;

  Vector4D res = posnwt[id] - H*pos[id]; //measurement residual
  Matrix4x4 S = H*P[id]*H.transpose() + R;
  MatrixNxM<7,4> K = P[id] * H.transpose() * S.inverse();


  pos[id] = pos[id] + K * res;

  // cheat and normalize;
  quatlength = 0;
  for (int i = 0; i < 4; i++)
    quatlength += pos[id](i)*pos[id](i);
  quatlength = sqrt(quatlength);
  for (int i = 0; i < 4; i++)
    pos[id](i) /= quatlength;

  P[id] = (MatrixNxN<7>::identity() - K*H) * P[id];

  if (DEBUG)
  {
    printf("** AFTER UPDATE **\n");
    printf("matrix R:\n");
    printmatrix(R);
    printf("matrix K:\n");
    printmatrix(K);
    printf("matrix P:\n");
    printmatrix(P[id]);
  }

// **************





}





Quaternion IMUProcessor::getPosRawAcc()
{
  return posrawacc[id];
}

Quaternion IMUProcessor::getPosRawGyr()
{
  return posrawgyr[id];
}

Quaternion IMUProcessor::getPosNewton()
{
  return posnwt[id];
}
Quaternion IMUProcessor::getPosRawAccError()
{
  return posrawaccerr[id];
}

Quaternion IMUProcessor::getPosRawGyrError()
{
  return posrawgyrerr[id];
}

Quaternion IMUProcessor::getPosNewtonError()
{
  return posnwterr[id];
}

VectorND<7> IMUProcessor::getPosKalman()
{
  return pos[id];
}

